import * as tf from '@tensorflow/tfjs';

function positionalEncoding(sequenceLength, dModel) {
  const position = tf.range(0, sequenceLength, 1).expandDims(1);
  const i = tf.range(0, dModel, 2).expandDims(0);
  
  const angleRates = tf.div(1, tf.pow(10000, tf.div(i, dModel)));
  const angleRads = tf.matMul(position, angleRates);
  
  const pe = tf.concat([
    tf.sin(angleRads),
    tf.cos(angleRads)
  ], 1);
  
  return pe.expandDims(0);
}

function transformerEncoderLayer(dModel, numHeads) {
  const input = tf.input({shape: [null, dModel]});
  
  // Multi-head attention
  const attention = tf.layers.multiHeadAttention({
    numHeads: numHeads,
    keyDim: dModel / numHeads
  }).apply([input, input, input]);
  
  // Add & Norm
  const addNorm1 = tf.layers.layerNormalization().apply(
    tf.layers.add().apply([input, attention])
  );
  
  // Feed forward
  const ff = tf.layers.dense({units: dModel * 4, activation: 'relu'}).apply(addNorm1);
  const ffOut = tf.layers.dense({units: dModel}).apply(ff);
  
  // Add & Norm
  const output = tf.layers.layerNormalization().apply(
    tf.layers.add().apply([addNorm1, ffOut])
  );
  
  return tf.model({inputs: input, outputs: output});
}

export function Transformer(numClasses, sequenceLength, dModel = 64, numHeads = 4, learningRate = 0.001) {
  const inputs = tf.input({shape: [sequenceLength, 1]});
  
  // Embedding
  const embedding = tf.layers.dense({units: dModel}).apply(inputs);
  
  // Positional encoding
  const pe = positionalEncoding(sequenceLength, dModel);
  const encoded = tf.layers.add().apply([embedding, pe]);
  
  // Transformer encoder
  const transformer = transformerEncoderLayer(dModel, numHeads).apply(encoded);
  
  // Global average pooling and output
  const pooled = tf.layers.globalAveragePooling1d().apply(transformer);
  const outputs = tf.layers.dense({
    units: numClasses,
    activation: 'softmax'
  }).apply(pooled);
  
  const model = tf.model({inputs, outputs});
  
  // Compile the model
  model.compile({
    optimizer: tf.train.adam(learningRate),
    loss: 'categoricalCrossentropy',
    metrics: ['accuracy']
  });
  
  return model;
}